#include "onnx_run.h"


YoloDetect::YoloDetect(const std::string &onnx_path, const float &conf_threshold) {
    this->confthreshold = conf_threshold;

    // 动态宽高时使用
    int dynamic_w = 640;
    int dynamic_h = 640;

    //environment （设置为VERBOSE（ORT_LOGGING_LEVEL_VERBOSE）时，方便控制台输出时看到是使用了cpu还是gpu执行）;
//    Ort::Env env = Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "ONNX_DETECTION");
    this->env = Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "ONNX_DETECTION");
//    Ort::SessionOptions sessionOptions;
    // 使用1个线程执行op,若想提升速度，增加线程数;
    this->sessionOptions.SetIntraOpNumThreads(1);
    // ORT_ENABLE_ALL: 启用所有可能的优化;
    this->sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

//    std::vector<std::string> availableProviders = Ort::GetAvailableProviders();
//    std::vector<std::string> availableProviders = {"TensorrtExecutionProvider","CUDAExecutionProvider","CPUExecutionProvider"};

    // TODO cuda;
    OrtCUDAProviderOptions cudaOption;
    this->sessionOptions.AppendExecutionProvider_CUDA(cudaOption);


#ifdef _WIN32
    // L表示双字节，每个字符占用两字节;
//    const wchar_t* model_path = LR"(D:\my_program\study-c\bushu\yolo-onnxruntime-test\models\yolov5s.onnx)";
    // wstring UTF-16编码 表示双字节，每个字符占用两字节;
    std::wstring w_onnx_path(onnx_path.begin(), onnx_path.end());
    // c_str()生成一个const char *指针，指向字符串的首地址;
    const wchar_t *model_path = w_onnx_path.c_str();
#else
    const char* model_path = onnx_path;
#endif

//    // 初始化模型;
    this->session = Ort::Session(this->env, model_path, this->sessionOptions);
    // 打印模型层;
    Ort::AllocatorWithDefaultOptions allocator;

    // 获取输入信息;
    Ort::TypeInfo inputTypeInfo = this->session.GetInputTypeInfo(0);
    Ort::TypeInfo outputTypeInfo = this->session.GetOutputTypeInfo(0);
    // 网络需要的输入形状,int64_t固定占用8个字节，不受编译器和操作系统影响，为了考虑可移植性;
    // 1, 3, 640, 640
    this->inputtensor_shape = inputTypeInfo.GetTensorTypeAndShapeInfo().GetShape();
    auto outputtensor_shape = outputTypeInfo.GetTensorTypeAndShapeInfo().GetShape();

    std::cout << "Input shape: ";
    for (auto shape: this->inputtensor_shape)
        std::cout << shape << ",";
    std::cout << std::endl;

    std::cout << "Output shape: ";
    for (auto shape: outputtensor_shape)
        std::cout << shape << ",";
    std::cout << std::endl;

    // 检查动态宽高
    if (this->inputtensor_shape[2] == -1 && this->inputtensor_shape[3] == -1) {
        // 动态宽高使用输入w和h
        this->input_size2f = cv::Size2f((float) dynamic_w, (float) dynamic_h);
        std::cout << "Dynamic shape h,w " << std::endl;
    } else {
        // 输入网络图像的宽高
        this->input_size2f = cv::Size2f((float) this->inputtensor_shape[3], (float) this->inputtensor_shape[2]);
    }

    //获取输入name;
    this->input_names.push_back(this->session.GetInputName(0, allocator));
    //获取输出name;
    this->output_names.push_back(this->session.GetOutputName(0, allocator));

    std::cout << "Input name: " << this->input_names[0] << std::endl;
    std::cout << "Output name: " << this->output_names[0] << std::endl;
    std::cout << "=========================" << std::endl;
}

std::vector<Detection> YoloDetect::detect(cv::Mat &draw_img) {
    // resize的图像和归一化之后的图像
    cv::Mat resize_image, float_image;
    // 存储缩放比例，左边和上边 pad 的值
    std::map<std::string, float> m_resize;
    // 真正输入网络的形状
    std::vector<int64_t> inputnew_shape = this->inputtensor_shape;  // {1,3,640,640}
    cv::Size realimg_shape = draw_img.size();  // 1080 x 810

    // 检查动态batch
//    if (this->inputtensor_shape[0] == -1) {
//        int batch = 2;
//        inputnew_shape[0] = batch;
//    }

    // resize等比缩放  [1080 x 810 --> 640 x 640]
    this->ScaleResize(draw_img, resize_image, m_resize);

    inputnew_shape[2] = resize_image.rows;
    inputnew_shape[3] = resize_image.cols;

    // BGR --> RGB 标准化;
    this->ImgNormalize(resize_image, float_image);

    // 堆区创建一个数组
    float *blob = new float[float_image.cols * float_image.rows * float_image.channels()];

    // cols->w,rows->h
    cv::Size floatimage_size{float_image.cols, float_image.rows};

    // h,w,c --> c,h,w;
    std::vector<cv::Mat> chw(float_image.channels());
    for (int i = 0; i < float_image.channels(); ++i) {
        // blob是数组首地址
        // 数组指针地址变形为 floatImageSize大小，成为 Mat 对象，赋给容器的第一个元素。数组地址加一个通道(w*h)的大小，进行下一次赋值;
        chw[i] = cv::Mat(floatimage_size, CV_32FC1, blob + i * floatimage_size.width * floatimage_size.height);
    }
    // 用split将图像切分通道赋值给chw容器，同时也是赋值给blob数组
    cv::split(float_image, chw);

    // 输入
    // 输入大小
    size_t inputnew_size = inputnew_shape[0] * inputnew_shape[1] * inputnew_shape[2] * inputnew_shape[3];
    // 输入值
    std::vector<float> inputtensor_values(blob, blob + inputnew_size);
    // 输入存放值的容器
    std::vector<Ort::Value> input_tensor;
    // 数据存储的内存类型，存储在CPU的输入Tensor数据，不影响模型在GPU上跑
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator,
                                                             OrtMemType::OrtMemTypeDefault);
    // 数据存储内存类型，核心数据，输入大小的字节数，数据的形状，数据形状的维度
    input_tensor.push_back(Ort::Value::CreateTensor<float>(memory_info, inputtensor_values.data(), inputnew_size,
                                                           inputnew_shape.data(), inputnew_shape.size()));

    // 输入张量获取输出值
    // 第一个值可忽略，模型输入名称，模型输入，输入个数，模型输出名称，输出个数
    std::vector<Ort::Value> output_tensor = this->session.Run(Ort::RunOptions{nullptr},
                                                              this->input_names.data(),
                                                              input_tensor.data(),
                                                              input_tensor.size(),
                                                              this->output_names.data(),
                                                              1);


    // 输出处理：置信度阈值筛选，nms，坐标反算
    std::vector<Detection> result = this->ScaleNms(output_tensor, realimg_shape, m_resize);

    delete[] blob;


    return result;
}

// 网络输出，检测图片形状，缩放比例和pad
std::vector<Detection> YoloDetect::ScaleNms(const std::vector<Ort::Value> &output_tensor, const cv::Size &realimg_shape,
                                            const std::map<std::string, float> &m_resize) {
    std::vector<cv::Rect> boxes;  // 框
    std::vector<float> confs;  // 置信度
    std::vector<int> clas_ids;  // 类别索引

    // 输出值地址
    auto *rawOutput = output_tensor[0].GetTensorData<float>();
    // 输出形状  [1,25200,85]  [x1,y1,x2,y2，置信度，类别索引]
    std::vector<int64_t> output_shape = output_tensor[0].GetTensorTypeAndShapeInfo().GetShape();
    // 输出大小 1*25200*85=2142000
    size_t count = output_tensor[0].GetTensorTypeAndShapeInfo().GetElementCount();
    std::vector<float> output(rawOutput, rawOutput + count);

    // 模型类别数
    int num_classes = (int) output_shape[2] - 5;
    // 单个批次的数据大小
    int elementsInBatch = (int) (output_shape[1] * output_shape[2]);

    // 获取缩放比例和左边上边填充像素值
    float scale = (float) m_resize.find("scale")->second;
    float top_pad = (float) m_resize.find("top")->second;
    float left_pad = (float) m_resize.find("left")->second;

    // 循环每个输出的框直到循环完所有框
    for (auto it = output.begin(); it != output.begin() + elementsInBatch; it += output_shape[2]) {
        // 坐标置信度
        float cls_conf = it[4];

        // 框的置信度小于阈值直接跳过,因为之后过滤阈值还会乘一个obj_conf进行过滤，都小于1，值只会更小
        if (cls_conf < this->confthreshold) {
            continue;
        }

        // 将类别的值转换成Mat对象
        std::vector<float> m1(it + 5, it + 5 + num_classes);
        cv::Mat clas_mat = cv::Mat(m1, true);  // [1 x 80]

        // 通过函数求最大最小值和索引，得到类别置信度和类别索引
        double max_value;
        cv::Point max_index;
        cv::minMaxLoc(clas_mat, nullptr, &max_value, nullptr, &max_index);

        // 类别置信度和类别索引
        int clas_id = max_index.y;
        float obj_conf = (float) max_value;

        // 获取cx,cy,w,h
        float cx = it[0];
        float cy = it[1];
        float ow = it[2];
        float oh = it[3];

        // 坐标反算
        cv::Rect box;
        // std::round 四舍五入
        box.x = (int) std::round(((cx - ow / 2) - left_pad) / scale);
        box.y = (int) std::round(((cy - oh / 2) - top_pad) / scale);
        box.width = (int) std::round(ow / scale);
        box.height = (int) std::round(oh / scale);

        // 坐标置信度和类别置信度的乘积
        float confidence = cls_conf * obj_conf;

        // emplace_back 和 push_back 比省略了拷贝和移动元素的过程
        boxes.emplace_back(box);
        confs.emplace_back(confidence);
        clas_ids.emplace_back(clas_id);
    }

    // NMS获取符合条件的索引
    std::vector<int> indexes;
    cv::dnn::NMSBoxes(boxes, confs, this->confthreshold, 0.45f, indexes);

    // 创建结构体，结果存放在容器中
    std::vector<Detection> detections;
    for (int idx: indexes) {
        Detection det;
        det.box = cv::Rect(boxes[idx]);
        det.conf = confs[idx];
        det.clas_id = clas_ids[idx];
        detections.emplace_back(det);
    }
    return detections;
}

// 缩放图像加黑边
void YoloDetect::ScaleResize(const cv::Mat &image, cv::Mat &resize_image, std::map<std::string, float> &m_resize) {
    float h = (float) image.size().height;
    float w = (float) image.size().width;
    int dims = image.channels();
    float IMG_HEIGHT = this->input_size2f.height;
    float IMG_WIDTH = this->input_size2f.width;

    float scale = std::min(IMG_HEIGHT / h, IMG_WIDTH / w);  // 缩放比例
    m_resize.insert(std::pair<std::string, float>("scale", scale));

//    cv::resize(image, resize_image, cv::Size(0, 0), scale, scale, cv::INTER_LINEAR);  // 等比缩放
    cv::resize(image, resize_image, cv::Size(0, 0), scale, scale, cv::INTER_CUBIC);  // 等比缩放

    float r_h = (float) resize_image.size().height;  // 变形后的大小
    float r_w = (float) resize_image.size().width;

    int top = int((IMG_HEIGHT - r_h) / 2);
    int bottom = int(IMG_HEIGHT - r_h - (float) top);
    int left = int((IMG_WIDTH - r_w) / 2);
    int right = int(IMG_WIDTH - r_w - (float) left);

    m_resize.insert(std::pair<std::string, float>("top", (float) top));
    m_resize.insert(std::pair<std::string, float>("left", (float) left));

    copyMakeBorder(resize_image, resize_image, top, bottom, left, right, cv::BORDER_CONSTANT,
                   cv::Scalar(114, 114, 114));
}

// BGR --> RGB 标准化 除255
void YoloDetect::ImgNormalize(const cv::Mat &resize_image, cv::Mat &float_image) {
    cv::cvtColor(resize_image, float_image, cv::COLOR_BGR2RGB);
    // 改成float
    float_image.convertTo(float_image, CV_32FC3, 1 / 255.0);
}

// BGR --> RGB 标准化 减均值除方差
void YoloDetect::ImgNormalize_(const cv::Mat &resize_image, cv::Mat &float_image) {
    // 改成float
    resize_image.convertTo(resize_image, CV_32FC3);

    //数据处理 标准化
    std::vector<cv::Mat> channels, channel_p;
    cv::split(resize_image, channels);
    cv::Mat R, G, B;
    B = channels.at(0);
    G = channels.at(1);
    R = channels.at(2);

    B = (B / 255. - 0.406) / 0.225;
    G = (G / 255. - 0.456) / 0.224;
    R = (R / 255. - 0.485) / 0.229;

    channel_p.push_back(R);
    channel_p.push_back(G);
    channel_p.push_back(B);

    cv::merge(channel_p, float_image);
}

//=================================utils===============================
std::vector<cv::Scalar> SetColor(const size_t &num_classes) {
    int size = 3;  // 最多有 3^3=27 种颜色
    int *arr_b = new int[size]{0, 255, 128};
    int *arr_g = new int[size]{255, 128, 0};
    int *arr_r = new int[size]{128, 0, 255};

    std::vector<cv::Scalar> clas_colors;
    cv::Scalar color;
    // 循环类别总数
    for (int i = 0; i < num_classes; i++) {
        // 转成5进制数
        int c = i % size;
        int ab = (int) i / size;
        int b = ab % 10;
        int a = (int) ab / 10;

        color = cv::Scalar(arr_b[a], arr_g[b], arr_r[c]);

        clas_colors.emplace_back(color);
    }

    delete[] arr_b;
    delete[] arr_g;
    delete[] arr_r;

    return clas_colors;
}

// 画框
void DrawImg(cv::Mat &img, const std::vector<Detection> &reslult, const std::vector<std::string> &class_names) {
    auto colors = SetColor(class_names.size());
    for (const Detection &detection: reslult) {
        int clas_index = detection.clas_id;
        int conf = (int) std::round(detection.conf * 100);
        int x = detection.box.x;
        int y = detection.box.y;
        int h = detection.box.height;
        int w = detection.box.width;

        std::string label = class_names[clas_index] + ": 0." + std::to_string(conf);

        cv::Scalar color = colors[clas_index];
        // 画框
        cv::rectangle(img, detection.box, color, 3);
        // 写字
        cv::putText(img, label, cv::Point(x, y - 8), cv::LINE_AA, 1.2, color, 3);
    }
}

